import yaml
import os
from pathlib import Path
from typing import Dict, Any, List
import logging

class EdgeConfig:
    """边侧配置管理"""
    
    def __init__(self, config_path: str = "config.yaml"):
        self.config_path = config_path
        self.raw_config = self._load_config()
        
        # 边侧基础配置
        self.edge_node = self._get_edge_config()
        self.cloud_server = self._get_cloud_config()
        self.inference = self._get_inference_config()
        self.dataset = self._get_dataset_config()
        self.api = self._get_api_config()
        self.log = self._get_log_config()
        self.monitoring = self._get_monitoring_config()
        self.aggregation = self._get_aggregation_config()
        self.registration = self._get_registration_config()
        
        # 确保目录存在
        self.ensure_directories()
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            config_file = Path(self.config_path)
            if config_file.exists():
                with open(config_file, 'r', encoding='utf-8') as f:
                    return yaml.safe_load(f)
            else:
                # 返回默认配置
                return self._get_default_config()
        except Exception as e:
            logging.warning(f"加载配置文件失败: {e}，使用默认配置")
            return self._get_default_config()
    
    def _get_default_config(self) -> Dict[str, Any]:
        """获取默认配置"""
        return {
            'edge_node': {
                'id': 'edge-default',
                'name': '边侧节点-默认',
                'location': '默认位置',
                'service_ip': '0.0.0.0',
                'service_port': 8001,
                'cpu': '8核',
                'memory': '16GB',
                'gpu': 'RTX 3060',
                'bandwidth': '1Gbps',
                'service_type': '设备运检I'
            },
            'cloud_server': {
                'host': 'localhost',
                'port': 8000,
                'base_url': 'http://localhost:8000/api/v1'
            },
            'inference': {
                'model_name': 'yolov8n.pt',
                'confidence_threshold': 0.5,
                'check_interval': 30,
                'batch_size': 1,
                'auto_load_default': True
            },
            'dataset': {
                'temp_dir': 'temp',
                'datasets_dir': 'datasets',
                'confidence_dir': 'confidence'
            },
            'api': {
                'host': '0.0.0.0',
                'port': 8001,
                'timeout': 30
            },
            'log': {
                'level': 'INFO',
                'format': '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            },
            'monitoring': {
                'update_interval': 60,
                'enable_metrics': True
            },
            'aggregation': {
                'difficulty_thresholds': {
                    'very_easy': 0.9,
                    'easy': 0.7,
                    'medium': 0.5,
                    'hard': 0.3,
                    'very_hard': 0.0
                },
                'max_upload_samples': 1000,
                'upload_retry_times': 3
            },
            'registration': {
                'initial_retry_interval': 30,
                'max_retry_interval': 300,
                'heartbeat_interval': 60,
                'connection_timeout': 10
            }
        }
    
    def _get_edge_config(self):
        """获取边侧节点配置"""
        config = self.raw_config.get('edge_node', {})
        
        class EdgeNodeConfig:
            id = config.get('id', 'edge-default')
            name = config.get('name', '边侧节点-默认')
            location = config.get('location', '默认位置')
            service_ip = config.get('service_ip', '0.0.0.0')
            service_port = config.get('service_port', 8001)
            cpu = config.get('cpu', '8核')
            memory = config.get('memory', '16GB')
            gpu = config.get('gpu', 'RTX 3060')
            bandwidth = config.get('bandwidth', '1Gbps')
            service_type = config.get('service_type', '设备运检I')
        
        return EdgeNodeConfig()
    
    def _get_cloud_config(self):
        """获取云端服务配置"""
        config = self.raw_config.get('cloud_server', {})
        
        class CloudConfig:
            host = config.get('host', 'localhost')
            port = config.get('port', 8000)
            base_url = config.get('base_url', f"http://{config.get('host', 'localhost')}:{config.get('port', 8000)}/api/v1")
        
        return CloudConfig()
    
    def _get_inference_config(self):
        """获取推理配置"""
        config = self.raw_config.get('inference', {})
        
        class InferenceConfig:
            model_name = config.get('model_name', 'yolov8n.pt')
            confidence_threshold = config.get('confidence_threshold', 0.5)
            check_interval = config.get('check_interval', 30)
            batch_size = config.get('batch_size', 1)
            auto_load_default = config.get('auto_load_default', True)
        
        return InferenceConfig()
    
    def _get_dataset_config(self):
        """获取数据集配置"""
        config = self.raw_config.get('dataset', {})
        
        class DatasetConfig:
            temp_dir = config.get('temp_dir', 'temp')
            datasets_dir = config.get('datasets_dir', 'datasets')
            confidence_dir = config.get('confidence_dir', 'confidence')
            max_temp_images = config.get('max_temp_images', 1000)
            cleanup_interval = config.get('cleanup_interval', 3600)
            single_dataset_mode = config.get('single_dataset_mode', True)
        
        return DatasetConfig()
    
    def _get_api_config(self):
        """获取API配置"""
        config = self.raw_config.get('api', {})
        
        class APIConfig:
            host = config.get('host', '0.0.0.0')
            port = config.get('port', 8001)
            timeout = config.get('timeout', 30)
        
        return APIConfig()
    
    def _get_log_config(self):
        """获取日志配置"""
        config = self.raw_config.get('log', {})
        
        class LogConfig:
            level = config.get('level', 'INFO')
            format = config.get('format', '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        
        return LogConfig()
    
    def _get_monitoring_config(self):
        """获取监控配置"""
        config = self.raw_config.get('monitoring', {})
        
        class MonitoringConfig:
            update_interval = config.get('update_interval', 60)
            enable_metrics = config.get('enable_metrics', True)
            health_check_interval = config.get('health_check_interval', 30)
            max_failure_count = config.get('max_failure_count', 3)
            auto_recovery = config.get('auto_recovery', True)
        
        return MonitoringConfig()
    
    def _get_aggregation_config(self):
        """获取聚合配置"""
        config = self.raw_config.get('aggregation', {})
        
        class AggregationConfig:
            difficulty_thresholds = config.get('difficulty_thresholds', {
                'very_easy': 0.9,
                'easy': 0.7,
                'medium': 0.5,
                'hard': 0.3,
                'very_hard': 0.0
            })
            max_upload_samples = config.get('max_upload_samples', 1000)
            upload_retry_times = config.get('upload_retry_times', 3)
        
        return AggregationConfig()
    
    def _get_registration_config(self):
        """获取注册配置"""
        config = self.raw_config.get('registration', {})
        
        class RegistrationConfig:
            initial_retry_interval = config.get('initial_retry_interval', 30)
            max_retry_interval = config.get('max_retry_interval', 300)
            heartbeat_interval = config.get('heartbeat_interval', 60)
            connection_timeout = config.get('connection_timeout', 10)
        
        return RegistrationConfig()
    
    def ensure_directories(self):
        """确保必要目录存在"""
        directories = [
            Path(self.dataset.temp_dir),
            Path(self.dataset.datasets_dir),
            Path(self.dataset.datasets_dir) / 'images',
            Path(self.dataset.datasets_dir) / 'labels',
            Path(self.dataset.datasets_dir) / self.dataset.confidence_dir,
            Path("models"),
            Path("models") / "backup",
            Path("sample_uploads"),
            Path("logs")
        ]
        
        for directory in directories:
            Path(directory).mkdir(parents=True, exist_ok=True)

# 全局配置实例
config = EdgeConfig() 